import argparse


def parse_args(args=None):
    parser = argparse.ArgumentParser()

    parser.add_argument("--script", required=False, action="store_true")
    parser.add_argument("-d", "--debug", required=False, type=str, default=None)
    parser.add_argument(
        "-c",
        "--train_dataset",
        required=True,
        type=str,
        help="train dataset for train bert",
    )
    parser.add_argument(
        "-t",
        "--test_dataset",
        type=str,
        default=None,
        help="test set for evaluate train set",
    )
    parser.add_argument(
        "-v",
        "--vocab_path",
        required=True,
        type=str,
        help="built vocab model path with bert-vocab",
    )
    parser.add_argument(
        "-o", "--output_path", required=True, type=str, help="ex)output/bert.model"
    )

    parser.add_argument(
        "-hs",
        "--hidden",
        type=int,
        default=768,
        help="hidden size of transformer model",
    )
    parser.add_argument("-l", "--layers", type=int, default=12, help="number of layers")
    parser.add_argument(
        "-a", "--attn_heads", type=int, default=12, help="number of attention heads"
    )
    parser.add_argument(
        "-s", "--seq_len", type=int, default=128, help="maximum sequence len"
    )

    parser.add_argument(
        "-b", "--batch_size", type=int, default=64, help="number of batch_size"
    )
    parser.add_argument("-e", "--epochs", type=int, default=10, help="number of epochs")
    parser.add_argument(
        "-w", "--num_workers", type=int, default=0, help="dataloader worker size"
    )

    parser.add_argument(
        "--device", default=0, help="Device to use for training, str or int (CUDA only)"
    )
    parser.add_argument(
        "--log_freq", type=int, default=10, help="printing loss every n iter: setting n"
    )
    parser.add_argument(
        "--corpus_lines", type=int, default=None, help="total number of lines in corpus"
    )
    parser.add_argument(
        "--device_ids",
        nargs="+",
        default=None,
        help="Device ids, str or int (CUDA only)",
    )
    parser.add_argument(
        "--on_memory", type=bool, default=True, help="Loading on memory: true or false"
    )

    parser.add_argument("--lr", type=float, default=1e-3, help="learning rate of adam")
    parser.add_argument(
        "--adam_weight_decay", type=float, default=0.01, help="weight_decay of adam"
    )
    parser.add_argument(
        "--adam_beta1", type=float, default=0.9, help="adam first beta value"
    )
    parser.add_argument(
        "--adam_beta2", type=float, default=0.999, help="adam first beta value"
    )

    parsed_args = parser.parse_args(args)
    if isinstance(parsed_args.device, str) and parsed_args.device.isdigit():
        parsed_args.device = int(parsed_args.device)
    if isinstance(parsed_args.device_ids, str) and parsed_args.device_ids.isdigit():
        parsed_args.device_ids = int(parsed_args.device_ids)
    return parsed_args
